#pragma once

#include "D3D11Declarations.hpp"
#include "D3D11Exception.hpp"
#include "D3D11ShaderConstant.hpp"
#include "D3D11SamplerState.hpp"
#include "D3D11Texture.hpp"
#include "D3D11DeviceContext.hpp"
#include "RenderPass.hpp"




namespace Engine
{
	namespace Video
	{
		typedef ID3D11DomainShader     Domain;
		typedef ID3D11VertexShader     Vertex;
		typedef ID3D11PixelShader      Pixel;
		typedef ID3D11GeometryShader   Geometry;
		typedef ID3D11ComputeShader    Compute;




		template<typename Type>
		class D3D11Shader
		{
		public:

			typedef std::pair<const std::string, ConstantsBuffer::ptr>   ConstantsPair;
			typedef std::pair<const std::string, Texture::ptr>           TexturePair;
			typedef std::pair<const std::string, SamplerState::ptr>      SamplerPair;

			typedef std::vector<D3D11_SIGNATURE_PARAMETER_DESC>   SignatureDesc;

			typedef std::map<const std::string, uint32>          Slot;
			typedef std::map<std::string, D3D11ShaderConstant>   ConstantsMap;
			typedef std::map<std::string, uint32>                VariableSlotMap;

		private:

			COM_PTR(Type)               m_shader;
			ID3D10BlobPtr               m_compiledShader;
			ID3D11ShaderReflectionPtr   m_reflector;
			D3D11_SHADER_DESC           m_shaderDesc;

			ConstantsMap                m_constants;
			VariableSlotMap             m_textures;
			VariableSlotMap             m_samplers;

		public:

			/**
			 * Create an empty shader.
			 */
			D3D11Shader()
			{}

			/**
			 * Create a shader from a d3d shader interface and it's compiled code (will be stored
			 * until a call to optimize).
			 */
			D3D11Shader(const COM_PTR(Type)& shader, const ID3D10BlobPtr& compiledShader)
				: m_shader(shader)
				, m_compiledShader(compiledShader)
			{
				ID3D11ShaderReflection* reflector = NULL; 
				D3D11_API(D3DReflect(m_compiledShader->GetBufferPointer(), m_compiledShader->GetBufferSize(), IID_ID3D11ShaderReflection, (void**) &reflector));
				m_reflector = ID3D11ShaderReflectionPtr(reflector);

				D3D11_API(m_reflector->GetDesc(&m_shaderDesc));


				//
				// Enumerate constant buffers
				//

				for(uint32 i = 0; i < m_shaderDesc.ConstantBuffers; ++i)
				{
					ID3D11ShaderReflectionConstantBuffer* buffer = m_reflector->GetConstantBufferByIndex(i);
					D3D11_SHADER_BUFFER_DESC desc;
					D3D11_API(buffer->GetDesc(&desc));

					if(i == 0 && strcmp(desc.Name, "$Globals") != 0)
						BOOSTEXT_THROW(exception(format("The first constants buffer should be named '$Globals' but is named '%1%' instead") % desc.Name));

					// The globals buffer has no name (it's always at index 0
					m_constants.insert(std::make_pair(i == 0 ? desc.Name : "", D3D11ShaderConstant(buffer, desc, i)));
				}


				//
				// Enumerate Shader resources
				//

				for(uint32 i = 0; i < m_shaderDesc.BoundResources; ++i)
				{
					D3D11_SHADER_INPUT_BIND_DESC desc;
					D3D11_API(m_reflector->GetResourceBindingDesc(i, &desc));
					switch(desc.Type)
					{
					case D3D10_SIT_TEXTURE:
						if(desc.Dimension != D3D10_SRV_DIMENSION_TEXTURE2D)
							BOOSTEXT_THROW(exception("Currently, only 2-dimensional textures are supported"));

						m_textures.insert(std::make_pair(std::string(desc.Name), desc.BindPoint));
						break;
					case D3D10_SIT_SAMPLER:
						m_samplers.insert(std::make_pair(std::string(desc.Name), desc.BindPoint));
						break;
					case D3D10_SIT_CBUFFER:
						// We've added those buffers already
						continue;
					default:
						BOOSTEXT_THROW(exception(format("The shader resource %1% is not supported yet") % desc.Type));
					}
				}
			}



			/**
			 * Optimize memory usage by freeing all reflection data and the compiled shader's buffer.
			 */
			void optimize()
			{
				m_compiledShader.release();
				m_reflector.release();
			}




			/**
			 * Get the Input Signature of this shader.
			 */
			SignatureDesc inputSignature() const
			{
				SignatureDesc signature;
				signature.resize(m_shaderDesc.InputParameters);
				for(uint32 i = 0; i < m_shaderDesc.InputParameters; ++i)
					D3D11_API(m_reflector->GetInputParameterDesc(i, &signature[i]));

				return signature;
			}

			/**
			 * Get the Output Signature of this shader.
			 */
			SignatureDesc outputSignature() const
			{
				SignatureDesc signature;
				signature.resize(m_shaderDesc.OutputParameters);
				for(uint32 i = 0; i < m_shaderDesc.OutputParameters; ++i)
					D3D11_API(m_reflector->GetInputParameterDesc(i, &signature[i]));

				return signature;
			}



			/**
			 * Get the shader constant with he given name.
			 */
			D3D11ShaderConstant& constant(const std::string& name = "$Globals")
			{
				ConstantsMap::iterator i = m_constants.find(name);
				if(i == m_constants.end())
					BOOSTEXT_THROW(exception(format("There is no shader constant named '%1%'") % name));

				return i->second;
			}

			/**
			 * Get the sampler's slot with the given name.
			 */
			uint32 sampler(const std::string& name)
			{
				VariableSlotMap::const_iterator i = m_samplers.find(name);
				if(i == m_samplers.end())
					BOOSTEXT_THROW(exception(format("There is no sampler named '%1%'") % name));

				return i->second;
			}

			/**
			 * Get the texture's slot with the given name.
			 */
			uint32 texture(const std::string& name)
			{
				VariableSlotMap::const_iterator i = m_textures.find(name);
				if(i == m_textures.end())
					BOOSTEXT_THROW(exception(format("There is no texture named '%1%'") % name));

				return i->second;
			}



			/**
			 * Get the compiled shader bytecode.
			 */
			ID3D10BlobPtr compiled() const
			{
				return m_compiledShader;
			}

			/**
			 * Apply the shader and all parameters to the context.
			 */
			void apply(const ID3D11DevicePtr& device, D3D11DeviceContext& context, const RenderPass::shader& params)
			{
				// Set all constant buffers
				foreach(const ConstantsPair& p, params.constants)
				{
					// Verify that the given layout matches the shaders expectations
					D3D11ShaderConstant& constantConfig = constant(p.first);
					constantConfig.compare(p.second->reflect());

					D3D11ConstantsBuffer::ptr constantBuffer = boost::shared_polymorphic_cast<D3D11ConstantsBuffer>(p.second);
					ID3D11Buffer* buffers[] = { constantBuffer->apply(device, context).get() };
					context.setConstantBuffers<Type>(0, 1, buffers);
				}

				// Set all texture samplers
				foreach(const SamplerPair& p, params.samplers)
				{
					uint32 index = sampler(p.first);

					D3D11SamplerState::ptr sampler = boost::shared_polymorphic_cast<D3D11SamplerState>(p.second);
					ID3D11SamplerState* samplers[] = { sampler->apply().get() };
					context.setSamplers<Type>(index, 1, samplers);
				}

				// Set all textures
				foreach(const TexturePair& p, params.textures)
				{
					uint32 index = texture(p.first);

					D3D11Texture::ptr texture = boost::shared_polymorphic_cast<D3D11Texture>(p.second);
					ID3D11ShaderResourceView* textures[] = { texture->view().get() };
					context.setShaderResources<Type>(index, 1, textures);
				}

				// Set the shader itself
				context.setShader(m_shader, NULL, 0);
			}



			/**
			 * Swap this shader with another one.
			 */
			void swap(D3D11Shader<Type>& that)
			{
				std::swap(m_shader,         that.m_shader);
				std::swap(m_compiledShader, that.m_compiledShader);
				std::swap(m_reflector,      that.m_reflector);
				std::swap(m_shaderDesc,     that.m_shaderDesc);

				std::swap(m_constants,      that.m_constants);
				std::swap(m_textures,       that.m_textures);
				std::swap(m_samplers,       that.m_samplers);
			}

		private:

			static Format::Type fromType(D3D10_REGISTER_COMPONENT_TYPE type, uint32 mask)
			{
				switch(type)
				{
				case D3D10_REGISTER_COMPONENT_FLOAT32:
					switch(mask)
					{
					case 15: return Format::R32G32B32A32_FLOAT;
					case 8:  return Format::R32G32B32_FLOAT;
					case 3:  return Format::R32G32_FLOAT;
					case 1:  return Format::R32_FLOAT;
					default: BOOSTEXT_THROW(invalid_value_exception(mask));
					}
					break;
				case D3D10_REGISTER_COMPONENT_UINT32:
					switch(mask)
					{
					case 15: return Format::R32G32B32A32_UINT;
					case 8:  return Format::R32G32B32_UINT;
					case 3:  return Format::R32G32_UINT;
					case 1:  return Format::R32_UINT;
					default: BOOSTEXT_THROW(invalid_value_exception(mask));
					}
					break;
				case D3D10_REGISTER_COMPONENT_SINT32:
					switch(mask)
					{
					case 15: return Format::R32G32B32A32_SINT;
					case 8:  return Format::R32G32B32_SINT;
					case 3:  return Format::R32G32_SINT;
					case 1:  return Format::R32_SINT;
					default: BOOSTEXT_THROW(invalid_value_exception(mask));
					}
					break;
				default: BOOSTEXT_THROW(invalid_value_exception(type));
				}
			}
		};
		///////////////////////////////////////////////////////////////////////////////////////////




		typedef D3D11Shader<Vertex>     D3D11VertexShader;
		typedef D3D11Shader<Pixel>      D3D11PixelShader;
		typedef D3D11Shader<Compute>    D3D11ComputeShader;
		typedef D3D11Shader<Geometry>   D3D11GeometryShader;
		typedef D3D11Shader<Domain>     D3D11DomainShader;




		inline ID3D10BlobPtr compile(const QByteArray& data, const std::string& name, const std::string& entryPoint, const std::string& profile, uint32 flags)
		{
			ID3D10Blob* shader;
			ID3D10Blob* errors;
			HRESULT hr = D3DCompile(data.constData(), data.size(), name.c_str(), NULL, NULL, entryPoint.c_str(), profile.c_str(), flags, NULL, &shader, &errors);
			if(hr == E_FAIL)
			{
				const char* error = (const char*)(errors->GetBufferPointer());
				BOOSTEXT_THROW(exception(format("Unable to compile shader using vs_4_0 profile: %1%") % (errors ? error : "Unknown problem")));
			}
			else
			{
				D3D11_API_F(hr, D3DCompile(data.constData(), data.size(), name.c_str(), NULL, NULL, entryPoint.c_str(), profile.c_str(), flags, NULL, &shader, &errors));
			}

			return ID3D10BlobPtr(shader);
		}
		///////////////////////////////////////////////////////////////////////////////////////////

		template <typename Type> inline
		ID3D10BlobPtr compileShader(const QByteArray& data, const std::string& entryPoint = "main");

		template <> inline
		ID3D10BlobPtr compileShader<Vertex>(const QByteArray& data, const std::string& entryPoint)
		{
			return compile(data, "vs", entryPoint, "vs_4_0", D3D10_SHADER_DEBUG | D3D10_SHADER_SKIP_OPTIMIZATION);
		}
		///////////////////////////////////////////////////////////////////////////////////////////

		template <> inline
		ID3D10BlobPtr compileShader<Pixel>(const QByteArray& data, const std::string& entryPoint)
		{
			return compile(data, "ps", entryPoint, "ps_4_0", D3D10_SHADER_DEBUG | D3D10_SHADER_SKIP_OPTIMIZATION);
		}
		///////////////////////////////////////////////////////////////////////////////////////////




		template <typename Type> inline
		D3D11Shader<Type> createShader(const ID3D11DevicePtr& device, const ID3D10BlobPtr& data);

		template <> inline
		D3D11Shader<Vertex> createShader<Vertex>(const ID3D11DevicePtr& device, const ID3D10BlobPtr& data)
		{
			ID3D11VertexShader* shader;
			D3D11_API(device->CreateVertexShader(data->GetBufferPointer(), data->GetBufferSize(), NULL, &shader));
			return D3D11VertexShader(ID3D11VertexShaderPtr(shader), data);
		}
		///////////////////////////////////////////////////////////////////////////////////////////

		template <> inline
		D3D11PixelShader createShader<Pixel>(const ID3D11DevicePtr& device, const ID3D10BlobPtr& data)
		{
			ID3D11PixelShader* shader;
			D3D11_API(device->CreatePixelShader(data->GetBufferPointer(), data->GetBufferSize(), NULL, &shader));
			return D3D11PixelShader(ID3D11PixelShaderPtr(shader), data);
		}
		///////////////////////////////////////////////////////////////////////////////////////////
	}
	///////////////////////////////////////////////////////////////////////////////////////////////
}
///////////////////////////////////////////////////////////////////////////////////////////////////




namespace std
{
	template <typename Type>
	void swap(Engine::Video::D3D11Shader<Type>& lhs, Engine::Video::D3D11Shader<Type>& rhs)
	{
		lhs.swap(rhs);
	}
}
///////////////////////////////////////////////////////////////////////////////////////////////////
